#!/usr/bin/env python3
import os
import cv2
import numpy as np
import time
import threading
import queue
from dataclasses import dataclass
from typing import Optional, Tuple
import glob
import traceback
# from display import PositionVisualizer
# from display_top import OptimizedPoseVisualizer
from opt_visualizer import OptimizedPoseVisualizer,OptimizedPositionVisualizer
from image_provider import ImageProvider
from seg_pose import OptimizedSegPose
from compute_position import ComputePosition  
from utils import get_depth_image
from logger import setup_logger

logger = setup_logger()

@dataclass
class ImageData:
    """图像数据结构"""
    timestamp: str
    readable_timestamp: str
    color_img: np.ndarray
    depth_img: np.ndarray
    color_path: Optional[str] = None

@dataclass
class ModelResult:
    """模型结果数据结构"""
    timestamp: str
    readable_timestamp: str
    color_img: np.ndarray
    depth_img: np.ndarray
    seg_result: any
    pose_result: any
    color_path: Optional[str] = None
    process_time: float = 0.0
    queue_wait_time: float = 0.0  # 新增：队列等待时间

@dataclass
class PositionResult:
    """位置计算结果数据结构"""
    timestamp: str
    readable_timestamp: str
    pose: any
    vector: any
    color_path: Optional[str] = None
    process_time: float = 0.0
    queue_wait_time: float = 0.0  # 新增：队列等待时间

class PerformanceMonitor:
    """性能监控类"""
    
    def __init__(self):
        self.lock = threading.Lock()
        self.stats = {
            'model_times': [],
            'position_times': [],
            'viz_times': [],
            'queue_wait_times': {
                'model': [],
                'position': [],
                'viz': []
            },
            'errors': {
                'model': 0,
                'position': 0,
                'viz': 0
            }
        }
    
    def record_time(self, category: str, time_ms: float):
        """记录处理时间"""
        with self.lock:
            if category in self.stats:
                self.stats[category].append(time_ms)
    
    def record_queue_wait(self, stage: str, wait_time_ms: float):
        """记录队列等待时间"""
        with self.lock:
            if stage in self.stats['queue_wait_times']:
                self.stats['queue_wait_times'][stage].append(wait_time_ms)
    
    def record_error(self, category: str):
        """记录错误"""
        with self.lock:
            if category in self.stats['errors']:
                self.stats['errors'][category] += 1
    
    def get_summary(self):
        """获取统计摘要"""
        with self.lock:
            summary = {}
            for key, times in self.stats.items():
                if key == 'queue_wait_times':
                    summary[key] = {}
                    for stage, wait_times in times.items():
                        if wait_times:
                            summary[key][stage] = {
                                'avg': sum(wait_times) / len(wait_times),
                                'max': max(wait_times),
                                'min': min(wait_times),
                                'count': len(wait_times)
                            }
                elif key == 'errors':
                    summary[key] = times
                elif times:  # 处理时间列表
                    summary[key] = {
                        'avg': sum(times) / len(times),
                        'max': max(times),
                        'min': min(times),
                        'count': len(times)
                    }
            return summary

class DiagnosticMultiThreadProcessor:
    """带诊断功能的多线程处理器"""
    
    def __init__(self, st_pose, st_vector, compute_position_instance, max_queue_size: int = 10):
        self.st_pose = st_pose
        self.st_vector = st_vector
        
        # 队列
        self.image_queue = queue.Queue(maxsize=max_queue_size)
        self.model_queue = queue.Queue(maxsize=max_queue_size)
        self.position_queue = queue.Queue(maxsize=max_queue_size)
        
        # 处理器 - 使用传入的已初始化实例
        self.seg_pose = OptimizedSegPose()
        self.compute_position = compute_position_instance  # 使用已经设置好标准位置的实例
        self.position_visualizer = OptimizedPositionVisualizer()
        self.pose_visualizer = OptimizedPoseVisualizer()
        
        # 性能监控
        self.performance_monitor = PerformanceMonitor()
        
        # 线程控制
        self.stop_event = threading.Event()
        self.threads = []
        
        # 统计信息
        self.stats = {
            'images_processed': 0,
            'models_processed': 0,
            'positions_processed': 0,
            'visualizations_processed': 0,
            'start_time': time.time(),
            'last_stats_time': time.time()
        }
        self.stats_lock = threading.Lock()
    
    def start(self):
        """启动所有线程"""
        # 模型处理线程
        model_thread = threading.Thread(target=self._model_worker, name="ModelWorker")
        model_thread.start()
        self.threads.append(model_thread)
        
        # 位置计算线程
        position_thread = threading.Thread(target=self._position_worker, name="PositionWorker")
        position_thread.start()
        self.threads.append(position_thread)
        
        # 可视化线程
        viz_thread = threading.Thread(target=self._visualization_worker, name="VizWorker")
        viz_thread.start()
        self.threads.append(viz_thread)
        
        # 性能监控线程
        monitor_thread = threading.Thread(target=self._performance_monitor_worker, name="MonitorWorker")
        monitor_thread.start()
        self.threads.append(monitor_thread)
        
        logger.info("所有处理线程已启动")
    
    def stop(self):
        """停止所有线程"""
        logger.info("正在停止所有线程...")
        self.stop_event.set()
        
        for thread in self.threads:
            thread.join(timeout=5.0)
            if thread.is_alive():
                logger.warning(f"线程 {thread.name} 未能在超时时间内停止")
        
        logger.info("所有线程已停止")
    
    def add_image(self, image_data: ImageData) -> bool:
        """添加图像到处理队列"""
        try:
            self.image_queue.put(image_data, timeout=0.1)
            with self.stats_lock:
                self.stats['images_processed'] += 1
            return True
        except queue.Full:
            logger.warning("图像队列已满，跳过当前图像")
            return False
    
    def _model_worker(self):
        """模型处理工作线程"""
        logger.info("模型处理线程已启动")
        
        while not self.stop_event.is_set():
            try:
                # 记录队列等待开始时间
                wait_start = time.time()
                image_data = self.image_queue.get(timeout=1.0)
                wait_time = (time.time() - wait_start) * 1000
                
                # 记录队列等待时间
                self.performance_monitor.record_queue_wait('model', wait_time)
                
                start_time = time.time()
                seg_result, pose_result = self.seg_pose.get_result(image_data.color_img)
                process_time = (time.time() - start_time) * 1000
                
                # 记录处理时间
                self.performance_monitor.record_time('model_times', process_time)
                
                model_result = ModelResult(
                    timestamp=image_data.timestamp,
                    readable_timestamp=image_data.readable_timestamp,
                    color_img=image_data.color_img,
                    depth_img=image_data.depth_img,
                    seg_result=seg_result,
                    pose_result=pose_result,
                    color_path=image_data.color_path,
                    process_time=process_time,
                    queue_wait_time=wait_time
                )
                
                self.model_queue.put(model_result, timeout=1.0)
                
                with self.stats_lock:
                    self.stats['models_processed'] += 1
                
                logger.debug(f"模型处理完成: {image_data.readable_timestamp}, "
                           f"处理耗时: {process_time:.2f}ms, 队列等待: {wait_time:.2f}ms")
                
            except queue.Empty:
                continue
            except queue.Full:
                logger.warning("模型结果队列已满")
            except Exception as e:
                logger.error(f"模型处理错误: {e}")
                logger.error(f"错误详情: {traceback.format_exc()}")
                self.performance_monitor.record_error('model')
        
        logger.info("模型处理线程已退出")
    
    def _position_worker(self):
        """位置计算工作线程 - 增强诊断"""
        logger.info("位置计算线程已启动")
        
        while not self.stop_event.is_set():
            try:
                # 记录队列等待开始时间
                wait_start = time.time()
                model_result = self.model_queue.get(timeout=1.0)
                wait_time = (time.time() - wait_start) * 1000
                
                # 记录队列等待时间
                self.performance_monitor.record_queue_wait('position', wait_time)
                
                logger.debug(f"位置计算开始: {model_result.readable_timestamp}")
                
                start_time = time.time()
                
                # 详细记录位置计算的各个步骤
                step1_start = time.time()
                pose, vector = self.compute_position.process_image_pair(
                    model_result.readable_timestamp,
                    model_result.color_img,
                    model_result.depth_img,
                    model_result.seg_result,
                    model_result.pose_result
                )
                step1_time = (time.time() - step1_start) * 1000
                
                process_time = (time.time() - start_time) * 1000
                
                # 记录处理时间
                self.performance_monitor.record_time('position_times', process_time)
                
                position_result = PositionResult(
                    timestamp=model_result.timestamp,
                    readable_timestamp=model_result.readable_timestamp,
                    pose=pose,
                    vector=vector,
                    color_path=model_result.color_path,
                    process_time=process_time,
                    queue_wait_time=wait_time
                )
                
                self.position_queue.put(position_result, timeout=1.0)
                
                with self.stats_lock:
                    self.stats['positions_processed'] += 1
                
                logger.info(f"位置计算完成: {model_result.readable_timestamp}, "
                          f"处理耗时: {process_time:.2f}ms, 队列等待: {wait_time:.2f}ms, "
                          f"内部计算: {step1_time:.2f}ms")
                logger.info(f"Pose: {pose}, Vector: {vector}")
                
            except queue.Empty:
                continue
            except queue.Full:
                logger.warning("位置结果队列已满")
            except Exception as e:
                logger.error(f"位置计算错误: {e}")
                logger.error(f"错误详情: {traceback.format_exc()}")
                self.performance_monitor.record_error('position')
        
        logger.info("位置计算线程已退出")
    
    def _visualization_worker(self):
        """可视化工作线程 - 增强诊断"""
        logger.info("可视化线程已启动")
        
        while not self.stop_event.is_set():
            try:
                # 记录队列等待开始时间
                wait_start = time.time()
                position_result = self.position_queue.get(timeout=1.0)
                wait_time = (time.time() - wait_start) * 1000
                
                # 记录队列等待时间
                self.performance_monitor.record_queue_wait('viz', wait_time)
                
                start_time = time.time()
                
                if position_result.color_path:
                    viz_start = time.time()
                    self.pose_visualizer.visualize_pose(
                        position_result.pose, 
                        self.st_pose, 
                        position_result.color_path
                    )
                    viz1_time = (time.time() - viz_start) * 1000
                    
                    viz2_start = time.time()
                    self.position_visualizer.visualize_position(
                        self.st_vector, 
                        position_result.vector, 
                        position_result.color_path
                    )
                    viz2_time = (time.time() - viz2_start) * 1000
                    
                    logger.debug(f"可视化详情 - 姿态: {viz1_time:.2f}ms, 位置: {viz2_time:.2f}ms")
                
                process_time = (time.time() - start_time) * 1000
                
                # 记录处理时间
                self.performance_monitor.record_time('viz_times', process_time)
                
                with self.stats_lock:
                    self.stats['visualizations_processed'] += 1
                
                logger.debug(f"可视化完成: {position_result.readable_timestamp}, "
                           f"处理耗时: {process_time:.2f}ms, 队列等待: {wait_time:.2f}ms")
                
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"可视化错误: {e}")
                logger.error(f"错误详情: {traceback.format_exc()}")
                self.performance_monitor.record_error('viz')
        
        logger.info("可视化线程已退出")
    
    def _performance_monitor_worker(self):
        """性能监控线程"""
        logger.info("性能监控线程已启动")
        
        while not self.stop_event.is_set():
            time.sleep(3.0)  # 每3秒输出一次详细统计
            
            with self.stats_lock:
                current_time = time.time()
                elapsed_time = current_time - self.stats['start_time']
                period_time = current_time - self.stats['last_stats_time']
                self.stats['last_stats_time'] = current_time
                
                if elapsed_time > 0:
                    fps = self.stats['images_processed'] / elapsed_time
                    
                    # 获取性能统计
                    perf_summary = self.performance_monitor.get_summary()
                    
                    logger.info("=== 详细性能诊断 ===")
                    logger.info(f"处理进度: 图像{self.stats['images_processed']} → 模型{self.stats['models_processed']} → 位置{self.stats['positions_processed']} → 可视化{self.stats['visualizations_processed']}")
                    logger.info(f"平均FPS: {fps:.2f}")
                    
                    # 处理时间统计
                    if 'model_times' in perf_summary:
                        model_stats = perf_summary['model_times']
                        logger.info(f"模型处理: 平均{model_stats['avg']:.2f}ms, 最大{model_stats['max']:.2f}ms, 最小{model_stats['min']:.2f}ms")
                    
                    if 'position_times' in perf_summary:
                        pos_stats = perf_summary['position_times']
                        logger.info(f"位置计算: 平均{pos_stats['avg']:.2f}ms, 最大{pos_stats['max']:.2f}ms, 最小{pos_stats['min']:.2f}ms")
                    
                    if 'viz_times' in perf_summary:
                        viz_stats = perf_summary['viz_times']
                        logger.info(f"可视化处理: 平均{viz_stats['avg']:.2f}ms, 最大{viz_stats['max']:.2f}ms, 最小{viz_stats['min']:.2f}ms")
                    
                    # 队列等待时间统计
                    if 'queue_wait_times' in perf_summary:
                        wait_stats = perf_summary['queue_wait_times']
                        for stage, stats in wait_stats.items():
                            if stats:
                                logger.info(f"{stage}队列等待: 平均{stats['avg']:.2f}ms, 最大{stats['max']:.2f}ms")
                    
                    # 错误统计
                    if 'errors' in perf_summary:
                        errors = perf_summary['errors']
                        total_errors = sum(errors.values())
                        if total_errors > 0:
                            logger.warning(f"错误统计: 模型{errors['model']}, 位置{errors['position']}, 可视化{errors['viz']}")
                    
                    # 队列状态
                    logger.info(f"队列状态: 图像{self.image_queue.qsize()}, 模型{self.model_queue.qsize()}, 位置{self.position_queue.qsize()}")
                    
                    # 瓶颈分析
                    model_backlog = self.model_queue.qsize()
                    position_backlog = self.position_queue.qsize()
                    
                    if model_backlog > 5:
                        logger.warning("⚠️  位置计算可能是瓶颈！模型结果队列积压过多")
                    elif position_backlog > 5:
                        logger.warning("⚠️  可视化可能是瓶颈！位置结果队列积压过多")
                    elif self.image_queue.qsize() == 0 and model_backlog == 0:
                        logger.info("✅ 流水线运行良好，无明显瓶颈")
                    
                    logger.info("==================")
        
        logger.info("性能监控线程已退出")

# 修改主函数以使用诊断处理器
def setup_standard_position():
    """设置标准位置 - 独立函数"""
    seg_pose = OptimizedSegPose()
    compute_position = ComputePosition()
    st_img = cv2.imread("./data/Color1/440.png")
    st_depth_img = get_depth_image("./data/Color1/440.png")

    logger.info("==========设置标准位置==========")
    seg_result, pose_result = seg_pose.get_result(st_img)
    if seg_result is None or pose_result is None:
        logger.error("==========没有检测到标准图片的seg结果或pose结果==========")
        return None, None, None

    st_pose, st_vector = compute_position.set_standard(st_img, st_depth_img, seg_result, pose_result)
    if st_pose is None or st_vector is None:
        logger.error("==========标准位置设定失败==========")
        return None, None, None
    else:
        logger.info("==========标准位置设定成功==========")
        logger.info(f"标准位置: {st_pose}")
        logger.info(f"标准向量: {st_vector}")
        return st_pose, st_vector, compute_position

def main():
    """主函数 - 使用诊断处理器"""
    # 设置标准位置
    st_pose, st_vector, compute_position = setup_standard_position()
    if st_pose is None or st_vector is None or compute_position is None:
        logger.error("标准位置设定失败，程序退出")
        return

    # 配置参数
    from_folder = False
    max_queue_size = 5  # 减小队列大小以便更容易发现瓶颈
    
    # 创建图像提供器
    # from processing_pipeline_cl import ImageProvider  # 使用原来的ImageProvider
    image_provider = ImageProvider(from_folder=from_folder)
    
    # 创建诊断处理器 - 传入已初始化的compute_position实例
    processor = DiagnosticMultiThreadProcessor(st_pose, st_vector, compute_position, max_queue_size)
    
    try:
        # 启动处理器
        processor.start()
        
        # 主线程负责读取图像并添加到队列
        logger.info("开始处理图像...")
        
        if from_folder:
            while True:
                image_data = image_provider.get_next_image()
                if image_data is None:
                    logger.info("所有图像读取完成")
                    break
                
                success = processor.add_image(image_data)
                if not success:
                    time.sleep(0.01)
        
        # 等待所有队列处理完成
        logger.info("等待队列处理完成...")
        while (not processor.image_queue.empty() or 
               not processor.model_queue.empty() or 
               not processor.position_queue.empty()):
            time.sleep(1.0)
            logger.info(f"等待中... 队列状态: 图像{processor.image_queue.qsize()}, 模型{processor.model_queue.qsize()}, 位置{processor.position_queue.qsize()}")
        
        logger.info("所有任务处理完成")
        
    except KeyboardInterrupt:
        logger.info("收到中断信号，正在关闭...")
    except Exception as e:
        logger.error(f"程序错误: {e}")
        logger.error(f"错误详情: {traceback.format_exc()}")
    finally:
        # 清理资源
        processor.stop()
        image_provider.close()
        
        # 输出最终详细统计
        with processor.stats_lock:
            total_time = time.time() - processor.stats['start_time']
            logger.info("==========程序结束==========")
            logger.info(f"总共处理图像: {processor.stats['images_processed']}")
            logger.info(f"总共模型处理: {processor.stats['models_processed']}")
            logger.info(f"总共位置计算: {processor.stats['positions_processed']}")
            logger.info(f"总共可视化: {processor.stats['visualizations_processed']}")
            if processor.stats['images_processed'] > 0:
                avg_fps = processor.stats['images_processed'] / total_time
                logger.info(f"平均帧率: {avg_fps:.2f} FPS")
            logger.info(f"总运行时间: {total_time:.2f} 秒")
            
            # 输出最终性能摘要
            perf_summary = processor.performance_monitor.get_summary()
            logger.info("==========性能摘要==========")
            for category, stats in perf_summary.items():
                if category not in ['queue_wait_times', 'errors'] and stats:
                    logger.info(f"{category}: 平均{stats['avg']:.2f}ms, 最大{stats['max']:.2f}ms")

if __name__ == "__main__":
    main()